"""
文档处理器
负责处理各种格式的文档
"""

from typing import List, Dict, Any, Optional
import os
import mimetypes
from pathlib import Path
import PyPDF2
import docx
from langchain.schema import Document

from app.core.config import settings


class DocumentProcessor:
    """文档处理器"""
    
    def __init__(self):
        self.supported_types = {
            'text/plain': self._process_txt,
            'application/pdf': self._process_pdf,
            'application/vnd.openxmlformats-officedocument.wordprocessingml.document': self._process_docx,
            'text/markdown': self._process_markdown
        }
    
    def process_file(self, file_path: str, metadata: Optional[Dict[str, Any]] = None) -> List[Document]:
        """处理文件并返回文档列表"""
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"文件不存在: {file_path}")
        
        # 获取文件类型
        mime_type, _ = mimetypes.guess_type(file_path)
        if not mime_type:
            # 根据文件扩展名判断
            ext = Path(file_path).suffix.lower()
            if ext == '.txt':
                mime_type = 'text/plain'
            elif ext == '.pdf':
                mime_type = 'application/pdf'
            elif ext == '.docx':
                mime_type = 'application/vnd.openxmlformats-officedocument.wordprocessingml.document'
            elif ext == '.md':
                mime_type = 'text/markdown'
            else:
                raise ValueError(f"不支持的文件类型: {ext}")
        
        # 检查是否支持该类型
        if mime_type not in self.supported_types:
            raise ValueError(f"不支持的文件类型: {mime_type}")
        
        # 处理文件
        processor = self.supported_types[mime_type]
        content = processor(file_path)
        
        # 构建文档元数据
        file_metadata = {
            "source": file_path,
            "filename": os.path.basename(file_path),
            "file_type": mime_type,
            "file_size": os.path.getsize(file_path)
        }
        
        if metadata:
            file_metadata.update(metadata)
        
        # 创建文档对象
        document = Document(
            page_content=content,
            metadata=file_metadata
        )
        
        return [document]
    
    def _process_txt(self, file_path: str) -> str:
        """处理文本文件"""
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                return f.read()
        except UnicodeDecodeError:
            # 尝试其他编码
            with open(file_path, 'r', encoding='gbk') as f:
                return f.read()
    
    def _process_pdf(self, file_path: str) -> str:
        """处理PDF文件"""
        try:
            with open(file_path, 'rb') as f:
                pdf_reader = PyPDF2.PdfReader(f)
                text = ""
                for page in pdf_reader.pages:
                    text += page.extract_text() + "\n"
                return text
        except Exception as e:
            raise Exception(f"PDF处理失败: {str(e)}")
    
    def _process_docx(self, file_path: str) -> str:
        """处理Word文档"""
        try:
            doc = docx.Document(file_path)
            text = ""
            for paragraph in doc.paragraphs:
                text += paragraph.text + "\n"
            return text
        except Exception as e:
            raise Exception(f"Word文档处理失败: {str(e)}")
    
    def _process_markdown(self, file_path: str) -> str:
        """处理Markdown文件"""
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                return f.read()
        except UnicodeDecodeError:
            with open(file_path, 'r', encoding='gbk') as f:
                return f.read()
    
    def validate_file(self, file_path: str) -> Dict[str, Any]:
        """验证文件"""
        if not os.path.exists(file_path):
            return {
                "valid": False,
                "error": "文件不存在"
            }
        
        # 检查文件大小
        file_size = os.path.getsize(file_path)
        if file_size > settings.max_file_size:
            return {
                "valid": False,
                "error": f"文件大小超过限制 ({file_size} > {settings.max_file_size})"
            }
        
        # 检查文件类型
        mime_type, _ = mimetypes.guess_type(file_path)
        if not mime_type:
            ext = Path(file_path).suffix.lower()
            if ext in ['.txt', '.pdf', '.docx', '.md']:
                return {
                    "valid": True,
                    "file_type": ext,
                    "file_size": file_size
                }
            else:
                return {
                    "valid": False,
                    "error": f"不支持的文件类型: {ext}"
                }
        
        if mime_type in self.supported_types:
            return {
                "valid": True,
                "file_type": mime_type,
                "file_size": file_size
            }
        else:
            return {
                "valid": False,
                "error": f"不支持的文件类型: {mime_type}"
            }
    
    def get_supported_types(self) -> List[str]:
        """获取支持的文件类型"""
        return list(self.supported_types.keys())
    
    def extract_metadata(self, file_path: str) -> Dict[str, Any]:
        """提取文件元数据"""
        stat = os.stat(file_path)
        return {
            "filename": os.path.basename(file_path),
            "file_size": stat.st_size,
            "created_time": stat.st_ctime,
            "modified_time": stat.st_mtime,
            "file_type": mimetypes.guess_type(file_path)[0]
        }
